package com.dwk;

/**
 * ①创建R_B_Tree，定义颜色
 * ②创建RBNode
 * ③辅助方法定义：parentOf(node)，isRed(node)， isBlack(node) ，setRed(node)，setBlack(node)，inOrderPrint()
 * ④左旋方法定义：leftRotate(node)
 * ⑤右旋方法定义：rightRotate(node)
 * ⑥公开插入接口方法定义：insert(K key, V value);
 * ⑦内部插入接口方法定义：insert(RBNode node);
 * ⑧修正插入导致红黑树失衡的方法定义：insertFIxUp(RBNode node);
 * ⑨测试红黑树正确性
 */

public class R_B_Tree<K extends Comparable<K>,V> {
    private static final boolean RED =true;
    private static final boolean BLACK =false;

    /**树根的引用*/
    private RBNode root;


    public RBNode getRoot() {
        return root;
    }

    /**
     * 获取当前节点的父节点
     */

    /**
     * 节点是否为红色
     */
    private boolean isRed(RBNode node){
        if(node!=null){
            return node.color==RED;
        }
        return false;
    }

    /**
     * 节点是否为黑色
     */
    private boolean isBlack(RBNode node){
        if(node!=null){
            return node.color==BLACK;
        }
        return false;
    }

    /**
     * 设置节点为红色
     */
    private void setRed(RBNode node){
        if (node!=null){
            node.color=RED;
        }
    }

    /**
     * 设置节点为黑色
     */
    private void setBlack(RBNode node){
        if (node!=null){
            node.color=BLACK;
        }
    }

    /**
     * 中序打印二叉树
     */

    public void inOrderPrint(){
        inOrderPrint(this.root);
    }

    private void inOrderPrint(RBNode node){
        if (node != null) {
            inOrderPrint(node.left);
            System.out.println("key:"+node.key+",value"+node.value);
            inOrderPrint(node.right);
        }
    }


    /**
     * 公开的插入方法
     */
    public void insert(K key,V value){
        RBNode node=new RBNode();
        node.setKey(key);
        node.setValue(value);
        //新节点一定是红色
        node.setColor(RED);
        insert(node);
    }

    private void insert(RBNode node){
        //查找当前node的插入位置
        RBNode parent=null;
        RBNode x=this.root;

        while (x!=null){
            parent=x;//找到当前node的父节点
            //a > b 则返回 1，否则返回 -1 ，相等返回0
            int cmp=node.key.compareTo(x.key);
            if(cmp<0){
                x=x.left;//向左继续查找
            }else if(cmp==0) {
                parent.setValue(node.value);
                return;
            }else {
                x=x.right;//向右继续查找
            }
        }
        node.parent=parent;

        //判断node应该插入到parent的左边还是右边
        if(parent != null) {
            if(node.key.compareTo(parent.key) < 0) {
                parent.left = node;
            } else {
                parent.right = node;
            }
        } else {//为空树,定义node为root
            this.root = node;
        }
        //插入之后需要进行修复红黑树，让红黑树再次平衡。
        insertFixUp(node);
    }

    /**
     * 插入后修复红黑树平衡的方法
     *     |---情景1：红黑树为空树,将根节点染色为黑色
     *     |---情景2：插入节点的key已经存在，不需要处理，前面returen掉了
     *     |---情景3：插入节点的父节点为黑色，红黑树仍然平衡，不需要处理
     *
     *     情景4 需要咱们去处理
     *     |---情景4：插入节点的父节点为红色
     *          |---情景4.1：叔叔节点存在，并且为红色（父-叔 双红）,将爸爸叔叔节点染色为黑色，将爷爷节点染色为红色，并且以爷爷节点为当前节点进行下一步处理
     *          |---情景4.2：叔叔节点不存在，或者为黑色，父节点为爷爷节点的左子树
     *               |---情景4.2.1：插入节点为其父节点的左子节点（LL情况）,将爸爸节点染色为黑色，将爷爷节点染色为红色,然后以爷爷节点右旋处理
     *               |---情景4.2.2：插入节点为其父节点的右子节点（LR情况），以爸爸节点左旋，变成LL情况，然后指定爸爸节点为当前节点进行下一步处理（
     *          |---情景4.3：叔叔节点不存在，或者为黑色，父节点为爷爷节点的右子树
     *               |---情景4.3.1：插入节点为其父节点的右子节点（RR情况），将爸爸节点染色为黑色，将爷爷节点染色为红色，然后以爷爷节点左旋处理
     *               |---情景4.3.2：插入节点为其父节点的左子节点（RL情况），以爸爸节点右旋，变成RR情况，然后指定爸爸节点为当前节点进行下一步处理
     */
    private void insertFixUp(RBNode node) {
        RBNode parent = parentOf(node);
        RBNode gparent = parentOf(parent);
        //存在父节点且父节点为红色
        if(parent != null && isRed(parent)) {
            //父节点是红色的，那么一定存在爷爷节点
            //父节点为爷爷节点的左子树
            if(parent == gparent.left) {
                RBNode uncle = gparent.right;
                //4.1：叔叔节点存在，并且为红色（父-叔 双红）
                //将父和叔染色为黑色，再将爷爷染红，并将爷爷设置为当前节点，进入下一次循环判断
                if(uncle != null && isRed(uncle)) {
                    setBlack(parent);
                    setBlack(uncle);
                    setRed(gparent);
                    insertFixUp(gparent);
                    return;
                }

                //叔叔节点不存在，或者为黑色，父节点为爷爷节点的左子树
                if(uncle == null || isBlack(uncle)) {
                    //插入节点为其父节点的右子节点（LR情况）=>
                    //左旋（父节点），当前节点设置为父节点，进入下一次循环
                    if(node == parent.right) {
                        leftRotate(parent);
                        insertFixUp(parent);
                        return;
                    }

                    //插入节点为其父节点的左子节点（LL情况）=>
                    //变色（父节点变黑，爷爷节点变红），右旋爷爷节点
                    if(node == parent.left) {
                        setBlack(parent);
                        setRed(gparent);
                        rightRotate(gparent);
                    }
                }
            }else {//父节点为爷爷节点的右子树
                RBNode uncle = gparent.left;
                //4.1：叔叔节点存在，并且为红色（父-叔 双红）
                //将父和叔染色为黑色，再将爷爷染红，并将爷爷设置为当前节点，进入下一次循环判断
                if(uncle != null && isRed(uncle)) {
                    setBlack(parent);
                    setBlack(uncle);
                    setRed(gparent);
                    insertFixUp(gparent);
                    return;
                }

                //叔叔节点不存在，或者为黑色，父节点为爷爷节点的右子树
                if(uncle == null || isBlack(uncle)) {
                    //插入节点为其父节点的左子节点（RL情况）
                    //右旋（父节点）得到RR情况，当前节点设置为父节点，进入下一次循环
                    if(node == parent.left) {
                        rightRotate(parent);
                        insertFixUp(parent);
                        return;
                    }

                    //插入节点为其父节点的右子节点（RR情况）
                    //变色（父节点变黑，爷爷节点变红），右旋爷爷节点
                    if(node == parent.right) {
                        setBlack(parent);
                        setRed(gparent);
                        leftRotate(gparent);
                    }
                }
            }
        }

        setBlack(this.root);
    }


    /**
     * 左旋方法
     * 左旋示意图：左旋x节点
     *    p                   p
     *    |                   |
     *    x                   y
     *   / \         ---->   / \
     *  lx  y               x   ry
     *     / \             / \
     *    ly  ry          lx  ly
     *
     * 左旋做了几件事？
     * 1.将x的右子节点指向y的左子节点,并将y的左子节的父节点更新为x,
     * 2.当x的父节点不为空时，将x的父节点指向y，并将x父节点的子树指向y
     * 3.将x的父节点更新为y，将y的左节点更新为x
     */
    private void leftRotate(RBNode x){
        RBNode y=x.right;
        //将x的右子节点指向y的左子节点
        x.right=y.left;

        //将y的左子节的父节点更新为x
        if(y.left!=null){
            y.left.parent=x;
        }

        //当x的父节点不为空时，将x的父节点指向y，并将x父节点的子树指向y
        if(x.parent!=null){
            y.parent=x.parent;
            if(x==x.parent.left){
                x.parent.left=y;
            }else {
                x.parent.right=y;
            }
        }else {//说明x为根节点，此时更新y为根节点
            this.root=y;
            this.root.parent=null;
        }

        //将x的父节点更新为y，将y的左节点更新为x
        x.parent=y;
        y.left=x;
    }


    /**
     * 右旋方法
     * 右旋示意图：右旋y节点
     *
     *    p                       p
     *    |                       |
     *    y                       x
     *   / \          ---->      / \
     *  x   ry                  lx  y
     * / \                         / \
     *lx  ly                      ly  ry
     *
     * 右旋都做了几件事？
     * 1.将x的右子节点 赋值 给了 y 的左子节点，并且更新x的右子节点的父节点为 y
     * 2.将y的父节点（不为空时）指向x的父节点，更新x的父节点为y的父节点
     * 3.将x的右子节点指向y，更新y的父节点为x
     */
    private void rightRotate(RBNode y) {
        //1.将x的右子节点赋值给y的左子节点，并将y赋值给x右子节点的父节点（x右子节点非空时）
       RBNode x=y.left;
       y.left=x.right;
       if(x.right!=null){
           x.right.parent=y;
       }
       //将y的父节点（不为空时）指向x，更新x的父节点为y的父节点
        if(y.parent!=null){
            x.parent=y.parent;
            if(y.parent.left==y){
                y.parent.left=x;
            }else {
                y.parent.right=x;
            }
        }else {//y为根节点
            this.root=x;
        }
        //将x的右子节点指向y，更新y的父节点为x
        x.right=y;
        y.parent=x;
    }

    //返回父节点
    private RBNode parentOf(RBNode node){
        if(node!=null){
            return node.parent;
        }
        return null;
    }

    static class RBNode<K extends Comparable<K>,V>{
        private RBNode parent;
        private RBNode left;
        private RBNode right;
        private boolean color;
        private K key;
        private V value;

        public RBNode(RBNode parent, RBNode left, RBNode right, boolean color, K key, V value) {
            this.parent = parent;
            this.left = left;
            this.right = right;
            this.color = color;
            this.key = key;
            this.value = value;
        }

        public RBNode() {
        }

        public RBNode getParent() {
            return parent;
        }

        public void setParent(RBNode parent) {
            this.parent = parent;
        }

        public RBNode getLeft() {
            return left;
        }

        public void setLeft(RBNode left) {
            this.left = left;
        }

        public RBNode getRight() {
            return right;
        }

        public void setRight(RBNode right) {
            this.right = right;
        }

        public boolean isColor() {
            return color;
        }

        public void setColor(boolean color) {
            this.color = color;
        }

        public K getKey() {
            return key;
        }

        public void setKey(K key) {
            this.key = key;
        }

        public V getValue() {
            return value;
        }

        public void setValue(V value) {
            this.value = value;
        }
    }
}
